import onnxruntime
import cv2
import numpy as np

CLASSES = ['Meter']  # coco80类别，这里假设模型只识别'Meter'类别


def load_and_crop_image(image_path, crop_width=1920, crop_height=1056):
    """
    从指定路径加载图像并裁剪为指定大小。

    参数:
        image_path (str): 图像文件的路径。
        crop_width (int): 裁剪后的宽度，默认1280。
        crop_height (int): 裁剪后的高度，默认1024。

    返回:
        cropped_image (numpy.ndarray): 裁剪后的图像，若加载或裁剪失败则返回 None。
    """
    # 加载图像
    image = cv2.imread(image_path)

    # 检查图像是否加载成功
    if image is None:
        print("图像加载失败，请检查路径")
        return None

    # 获取图像中心坐标
    center_x, center_y = image.shape[1] // 2, image.shape[0] // 2

    # 计算裁剪区域的左上角和右下角坐标
    start_x = max(center_x - crop_width // 2, 0)
    start_y = max(center_y - crop_height // 2, 0)
    end_x = min(start_x + crop_width, image.shape[1])
    end_y = min(start_y + crop_height, image.shape[0])

    # 裁剪图像
    cropped_image = image[start_y:end_y, start_x:end_x]

    # 检查裁剪后的图像尺寸是否符合预期
    if cropped_image.shape[1] != crop_width or cropped_image.shape[0] != crop_height:
        print("裁剪尺寸不匹配，可能原图尺寸不足")
        return None

    return cropped_image


# NMS非极大值抑制
def nms(dets, thresh):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    areas = (y2 - y1 + 1) * (x2 - x1 + 1)
    scores = dets[:, 4]
    keep = []
    index = scores.argsort()[::-1]

    while index.size > 0:
        i = index[0]
        keep.append(i)
        x11 = np.maximum(x1[i], x1[index[1:]])
        y11 = np.maximum(y1[i], y1[index[1:]])
        x22 = np.minimum(x2[i], x2[index[1:]])
        y22 = np.minimum(y2[i], y2[index[1:]])
        w = np.maximum(0, x22 - x11 + 1)
        h = np.maximum(0, y22 - y11 + 1)
        overlaps = w * h
        ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
        idx = np.where(ious <= thresh)[0]
        index = index[idx + 1]

    return keep


# 转换格式：从xywh到xyxy
def xywh2xyxy(x):
    y = np.copy(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2
    y[:, 1] = x[:, 1] - x[:, 3] / 2
    y[:, 2] = x[:, 0] + x[:, 2] / 2
    y[:, 3] = x[:, 1] + x[:, 3] / 2
    return y


# 过滤掉无用的框
def filter_box(org_box, conf_thres, iou_thres):
    org_box = np.squeeze(org_box)
    conf = org_box[..., 4] > conf_thres
    box = org_box[conf == True]
    cls_cinf = box[..., 5:]
    cls = [int(np.argmax(cls_cinf[i])) for i in range(len(cls_cinf))]
    all_cls = list(set(cls))
    output = []

    for i in range(len(all_cls)):
        curr_cls = all_cls[i]
        curr_cls_box = []
        for j in range(len(cls)):
            if cls[j] == curr_cls:
                box[j][5] = curr_cls
                curr_cls_box.append(box[j][:6])

        curr_cls_box = np.array(curr_cls_box)
        curr_cls_box = xywh2xyxy(curr_cls_box)
        curr_out_box = nms(curr_cls_box, iou_thres)

        for k in curr_out_box:
            output.append(curr_cls_box[k])
    output = np.array(output)
    return output


# 画框函数
def draw(image, box_data):
    boxes = box_data[..., :4].astype(np.int32)
    scores = box_data[..., 4]
    classes = box_data[..., 5].astype(np.int32)

    for box, score, cl in zip(boxes, scores, classes):
        top, left, right, bottom = box
        cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
        cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score), (top, left), cv2.FONT_HERSHEY_SIMPLEX, 0.6,
                    (0, 0, 255), 2)
    return image


# 封装成函数，接收图像、置信度、IOU阈值和模型路径
def detect_objects(img, model_path, conf_thres=0.25, iou_thres=0.45, expected_size=(1920, 1056)):
    # 加载ONNX模型
    session = onnxruntime.InferenceSession(model_path)

    # 获取图像的高度和宽度
    h, w, _ = img.shape

    # 检查图像尺寸是否符合模型输入要求
    if (w, h) != expected_size:
        print(f"错误: 输入图像的尺寸 {w}x{h} 不符合模型要求的尺寸 {expected_size[0]}x{expected_size[1]}")
        return None

    # 预处理图片
    or_img = cv2.resize(img, expected_size)
    img = or_img[:, :, ::-1].transpose(2, 0, 1)  # BGR2RGB和HWC2CHW
    img = img.astype(np.float32)
    img /= 255.0
    img = np.expand_dims(img, axis=0)

    # 推断
    outputs = session.run(None, {session.get_inputs()[0].name: img})

    # 处理输出
    predictions = outputs[0]
    outbox = filter_box(predictions, conf_thres, iou_thres)

    # 输出结果
    if len(outbox) == 0:
        print('未找到目标物体。')
        return None
    else:
        print('找到目标物体。')
        return outbox


# 示例调用
if __name__ == "__main__":
    image_path = r'D:\Users\Administrator\Desktop\VoltWatcherDemo\stitched_result.jpg'
    model_path = r'D:\Users\Administrator\Desktop\VoltWatcherDemo\Meter1920x4096.onnx'
    image = load_and_crop_image(image_path)
    outbox = detect_objects(image, model_path, conf_thres=0.5, iou_thres=0.4)

    # 如果你希望再画一次返回的框，可以单独调用 `draw` 函数
    if outbox is not None:
        img = cv2.imread(image_path)
        img_with_boxes = draw(img, outbox)  # 再次绘制框
        cv2.imshow("Marked Image", img_with_boxes)  # 显示标记后的图像
        cv2.waitKey(0)
        cv2.destroyAllWindows()
